"""
embeddings.py
=============
Embedding Generation Utilities

Provides:
- Text embedding generation using various models
- Batch embedding processing
- Caching for efficiency

Author: Neo4j Tutorial
Date: January 2025
"""

import os
from typing import List, Optional, Dict, Any, Union
from abc import ABC, abstractmethod
import hashlib
import json
import numpy as np
from dotenv import load_dotenv

load_dotenv()


class BaseEmbeddingModel(ABC):
    """Abstract base class for embedding models."""
    
    @abstractmethod
    def embed_text(self, text: str) -> List[float]:
        """Generate embedding for a single text."""
        pass
    
    @abstractmethod
    def embed_texts(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for multiple texts."""
        pass
    
    @property
    @abstractmethod
    def dimension(self) -> int:
        """Return the embedding dimension."""
        pass


class SentenceTransformerEmbedding(BaseEmbeddingModel):
    """
    Embedding model using sentence-transformers library.
    
    Recommended models:
    - all-MiniLM-L6-v2 (384 dims, fast)
    - all-mpnet-base-v2 (768 dims, better quality)
    - all-distilroberta-v1 (768 dims)
    """
    
    def __init__(
        self,
        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
        device: str = "cpu"
    ):
        """
        Initialize sentence transformer model.
        
        Args:
            model_name: Model name or path
            device: Device to use (cpu/cuda)
        """
        self.model_name = model_name
        self.device = device
        self._model = None
        self._dimension = None
    
    @property
    def model(self):
        """Lazy load the model."""
        if self._model is None:
            try:
                from sentence_transformers import SentenceTransformer
                self._model = SentenceTransformer(self.model_name, device=self.device)
                self._dimension = self._model.get_sentence_embedding_dimension()
                print(f"✅ Loaded embedding model: {self.model_name} ({self._dimension} dims)")
            except ImportError:
                raise ImportError(
                    "sentence-transformers not installed. "
                    "Run: pip install sentence-transformers"
                )
        return self._model
    
    @property
    def dimension(self) -> int:
        """Return embedding dimension."""
        if self._dimension is None:
            _ = self.model  # Trigger lazy load
        return self._dimension
    
    def embed_text(self, text: str) -> List[float]:
        """
        Generate embedding for a single text.
        
        Args:
            text: Input text
            
        Returns:
            Embedding vector as list of floats
        """
        embedding = self.model.encode(text, convert_to_numpy=True)
        return embedding.tolist()
    
    def embed_texts(
        self,
        texts: List[str],
        batch_size: int = 32,
        show_progress: bool = False
    ) -> List[List[float]]:
        """
        Generate embeddings for multiple texts.
        
        Args:
            texts: List of input texts
            batch_size: Batch size for processing
            show_progress: Show progress bar
            
        Returns:
            List of embedding vectors
        """
        embeddings = self.model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=show_progress,
            convert_to_numpy=True
        )
        return embeddings.tolist()


class OpenAIEmbedding(BaseEmbeddingModel):
    """
    Embedding model using OpenAI API.
    
    Models:
    - text-embedding-3-small (1536 dims)
    - text-embedding-3-large (3072 dims)
    - text-embedding-ada-002 (1536 dims, legacy)
    """
    
    def __init__(
        self,
        model_name: str = "text-embedding-3-small",
        api_key: Optional[str] = None
    ):
        """
        Initialize OpenAI embedding model.
        
        Args:
            model_name: OpenAI model name
            api_key: API key (uses OPENAI_API_KEY env var if not provided)
        """
        self.model_name = model_name
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        self._client = None
        
        # Model dimensions
        self._dimensions = {
            "text-embedding-3-small": 1536,
            "text-embedding-3-large": 3072,
            "text-embedding-ada-002": 1536
        }
    
    @property
    def client(self):
        """Lazy load OpenAI client."""
        if self._client is None:
            try:
                from openai import OpenAI
                self._client = OpenAI(api_key=self.api_key)
                print(f"✅ Initialized OpenAI embedding model: {self.model_name}")
            except ImportError:
                raise ImportError(
                    "openai not installed. Run: pip install openai"
                )
        return self._client
    
    @property
    def dimension(self) -> int:
        """Return embedding dimension."""
        return self._dimensions.get(self.model_name, 1536)
    
    def embed_text(self, text: str) -> List[float]:
        """
        Generate embedding for a single text using OpenAI.
        
        Args:
            text: Input text
            
        Returns:
            Embedding vector
        """
        response = self.client.embeddings.create(
            model=self.model_name,
            input=text
        )
        return response.data[0].embedding
    
    def embed_texts(
        self,
        texts: List[str],
        batch_size: int = 100
    ) -> List[List[float]]:
        """
        Generate embeddings for multiple texts.
        
        Args:
            texts: List of input texts
            batch_size: Batch size (OpenAI limit)
            
        Returns:
            List of embedding vectors
        """
        all_embeddings = []
        
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            response = self.client.embeddings.create(
                model=self.model_name,
                input=batch
            )
            # Sort by index to maintain order
            sorted_data = sorted(response.data, key=lambda x: x.index)
            all_embeddings.extend([d.embedding for d in sorted_data])
        
        return all_embeddings


class EmbeddingGenerator:
    """
    High-level embedding generator with caching support.
    """
    
    def __init__(
        self,
        model_type: str = "sentence-transformer",
        model_name: Optional[str] = None,
        cache_enabled: bool = True,
        cache_dir: str = ".embedding_cache"
    ):
        """
        Initialize embedding generator.
        
        Args:
            model_type: 'sentence-transformer' or 'openai'
            model_name: Specific model name
            cache_enabled: Enable embedding caching
            cache_dir: Directory for cache files
        """
        self.cache_enabled = cache_enabled
        self.cache_dir = cache_dir
        self._cache: Dict[str, List[float]] = {}
        
        # Initialize model
        if model_type == "sentence-transformer":
            model_name = model_name or os.getenv(
                "EMBEDDING_MODEL",
                "sentence-transformers/all-MiniLM-L6-v2"
            )
            self.model = SentenceTransformerEmbedding(model_name)
        elif model_type == "openai":
            model_name = model_name or "text-embedding-3-small"
            self.model = OpenAIEmbedding(model_name)
        else:
            raise ValueError(f"Unknown model type: {model_type}")
    
    def _get_cache_key(self, text: str) -> str:
        """Generate cache key for text."""
        return hashlib.md5(text.encode()).hexdigest()
    
    @property
    def dimension(self) -> int:
        """Get embedding dimension."""
        return self.model.dimension
    
    def embed(self, text: str, use_cache: bool = True) -> List[float]:
        """
        Generate embedding for text with optional caching.
        
        Args:
            text: Input text
            use_cache: Use cache if available
            
        Returns:
            Embedding vector
        """
        if use_cache and self.cache_enabled:
            cache_key = self._get_cache_key(text)
            if cache_key in self._cache:
                return self._cache[cache_key]
        
        embedding = self.model.embed_text(text)
        
        if self.cache_enabled:
            self._cache[cache_key] = embedding
        
        return embedding
    
    def embed_batch(
        self,
        texts: List[str],
        use_cache: bool = True,
        show_progress: bool = False
    ) -> List[List[float]]:
        """
        Generate embeddings for multiple texts.
        
        Args:
            texts: List of input texts
            use_cache: Use cache where available
            show_progress: Show progress indicator
            
        Returns:
            List of embedding vectors
        """
        if not use_cache or not self.cache_enabled:
            return self.model.embed_texts(texts)
        
        # Check cache for existing embeddings
        results = [None] * len(texts)
        texts_to_embed = []
        indices_to_embed = []
        
        for i, text in enumerate(texts):
            cache_key = self._get_cache_key(text)
            if cache_key in self._cache:
                results[i] = self._cache[cache_key]
            else:
                texts_to_embed.append(text)
                indices_to_embed.append(i)
        
        # Generate embeddings for non-cached texts
        if texts_to_embed:
            if show_progress:
                print(f"   Generating {len(texts_to_embed)} new embeddings...")
            new_embeddings = self.model.embed_texts(texts_to_embed)
            
            for idx, embedding in zip(indices_to_embed, new_embeddings):
                results[idx] = embedding
                cache_key = self._get_cache_key(texts[idx])
                self._cache[cache_key] = embedding
        
        return results
    
    def clear_cache(self):
        """Clear the embedding cache."""
        self._cache.clear()
    
    def get_cache_stats(self) -> Dict[str, int]:
        """Get cache statistics."""
        return {
            "cached_embeddings": len(self._cache),
            "dimension": self.dimension
        }


def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
    """
    Calculate cosine similarity between two vectors.
    
    Args:
        vec1: First vector
        vec2: Second vector
        
    Returns:
        Cosine similarity score
    """
    a = np.array(vec1)
    b = np.array(vec2)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))


def euclidean_distance(vec1: List[float], vec2: List[float]) -> float:
    """
    Calculate Euclidean distance between two vectors.
    
    Args:
        vec1: First vector
        vec2: Second vector
        
    Returns:
        Euclidean distance
    """
    a = np.array(vec1)
    b = np.array(vec2)
    return float(np.linalg.norm(a - b))


# ============================================================
# TESTING
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Testing Embedding Generator")
    print("=" * 60)
    
    # Test with sentence transformer
    generator = EmbeddingGenerator(model_type="sentence-transformer")
    
    # Single embedding
    text = "Neo4j is a graph database management system."
    print(f"\nText: {text}")
    
    embedding = generator.embed(text)
    print(f"Embedding dimension: {len(embedding)}")
    print(f"First 5 values: {embedding[:5]}")
    
    # Batch embeddings
    texts = [
        "Graph databases store data as nodes and relationships.",
        "Neo4j uses Cypher as its query language.",
        "Knowledge graphs represent real-world entities."
    ]
    
    print(f"\nGenerating embeddings for {len(texts)} texts...")
    embeddings = generator.embed_batch(texts)
    
    # Calculate similarities
    print("\nSimilarity matrix:")
    for i, t1 in enumerate(texts):
        for j, t2 in enumerate(texts):
            sim = cosine_similarity(embeddings[i], embeddings[j])
            print(f"   [{i},{j}]: {sim:.4f}")
    
    # Cache stats
    print(f"\nCache stats: {generator.get_cache_stats()}")
    
    print("\n✅ Embedding test complete!")
